This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

Greptile Summary

This PR fixes a critical bug in ObliqueSGD where the momentum buffer was never actually persisted across optimizer steps — buf = torch.add(...) created a new tensor each step without updating state["momentum_buffer"], so the optimizer silently ran with zero effective momentum. The fix correctly uses torch.add(..., out=buf) to update the buffer in-place. Alongside this, the helper function _compute_riemannian_grad_and_update is split into a pure _compute_riemannian_grad and inline weight-decay/update logic, enabling both ObliqueSGD and ObliqueAdam to integrate with the existing WeightDecayMixin and support configurable weight-decay methods (decoupled, independent, l2).

• Bug fix is correct — the in-place out=buf write is the right approach; a new test (test_oblique_sgd_momentum_buffer_accumulates_across_steps) directly validates both-step buffer accumulation.
• Convergence threshold lowered 50 % → 40 % — since working momentum should improve (not hurt) convergence, this reduction is suspicious and may indicate the high weight_decay=0.1 + now-active momentum causes instability on this test setup, or the test was always flaky without a fixed seed.
• weight_decay_method not in defaults dict — the parameter is stored only as self.weight_decay_method, so it is not preserved by state_dict()/load_state_dict() and cannot vary per param group, unlike the other optimizer hyperparameters (lr, momentum, weight_decay).

Confidence Score: 4/5

• Safe to merge after confirming why the convergence accuracy threshold was lowered
• The core bug fix is correct and well-tested. The refactoring is clean and the WeightDecayMixin integration is consistent with the rest of the codebase. One P1 concern remains: the convergence threshold reduction from 50% to 40% is counter-intuitive after fixing a bug that should improve convergence, and warrants a brief explanation or investigation before merge. The missing weight_decay_method in defaults is a follow-up style concern that does not block correctness for current usage.
• tests/convergence/normalized_optimizer_test.py — the accuracy threshold reduction needs justification

Important Files Changed

Sequence Diagram

sequenceDiagram
    participant Caller
    participant ObliqueSGD/ObliqueAdam
    participant WeightDecayMixin
    participant _compute_riemannian_grad

    Caller->>ObliqueSGD/ObliqueAdam: step()
    ObliqueSGD/ObliqueAdam->>ObliqueSGD/ObliqueAdam: update momentum buffer in-place<br/>torch.add(grad, buf, alpha=mom, out=buf)
    ObliqueSGD/ObliqueAdam->>_compute_riemannian_grad: (param, buf/norm_grad, dim)
    _compute_riemannian_grad-->>ObliqueSGD/ObliqueAdam: riem_grad (tangent-space projected)
    ObliqueSGD/ObliqueAdam->>WeightDecayMixin: _apply_weight_decay_inplace(param, riem_grad, lr, wd)
    Note over WeightDecayMixin: decoupled: param *= (1 - lr*wd)<br/>independent: param *= (1 - wd)<br/>l2: riem_grad += wd*param
    WeightDecayMixin-->>ObliqueSGD/ObliqueAdam: param or riem_grad updated in-place
    ObliqueSGD/ObliqueAdam->>ObliqueSGD/ObliqueAdam: param.add_(riem_grad, alpha=-lr)
    ObliqueSGD/ObliqueAdam->>ObliqueSGD/ObliqueAdam: normalize(param) — retract to manifold

_{Reviews (3): Last reviewed commit: "addressed failing test" | Re-trigger Greptile}

/ok to test ec6b492

/ok to test 7091f0c